/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2024 All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include <iostream>
#include <gtest/gtest.h>
#include "ascir_ops.h"
#include "ascend_ir/expression.h"
#include "graph/utils/cg_utils.h"
namespace ge {
namespace ascir {
namespace cg {
using Graph = ge::AscGraph;
using ge::expression::Expression;
using ge::expression::ExpressionPtr;
TEST(CgUtils, SetGetContextOk) {
  Graph graph("test_graph");
  auto A = Expression::CreateExpression("A");
  auto B = Expression::CreateExpression("B");
  auto C = Expression::CreateExpression("C");
  auto a = graph.CreateAxis("a", A);
  auto b = graph.CreateAxis("b", B);
  auto c = graph.CreateAxis("c", C);
  auto ctx = CgContext::GetSharedThreadLocalContext();
  ASSERT_EQ(ctx, nullptr);
  auto ctx_obj = std::make_shared<CgContext>();
  CgContext::SetThreadLocalContext(ctx_obj);
  ctx = CgContext::GetSharedThreadLocalContext();
  ASSERT_NE(ctx, nullptr);
  ctx->SetLoopAxes({a, b, c});
  ASSERT_EQ(ctx->GetLoopAxes().size(), 3);
  ctx->SetBlockLoopEnd(a.id);
  ASSERT_EQ(ctx->GetBlockLoopEnd(), a.id);
  ctx->SetVectorizedLoopEnd(c.id);
  ASSERT_EQ(ctx->GetVectorizedLoopEnd(), c.id);
  ctx->SetLoopEnd(c.id);
  ASSERT_EQ(ctx->GetLoopEnd(), c.id);
}

TEST(CgUtils, LoopGuardContextOk) {
  Graph graph("test_graph");
  auto A = Expression::CreateExpression("A");
  auto B = Expression::CreateExpression("B");
  auto C = Expression::CreateExpression("C");
  auto a = graph.CreateAxis("a", A);
  auto b = graph.CreateAxis("b", B);
  auto c = graph.CreateAxis("c", C);

  int64_t count = 0;
  ASSERT_EQ(CgContext::GetThreadLocalContext(), nullptr);
  ASSERT_EQ(CgContext::GetSharedThreadLocalContext(), nullptr);
  LOOP(a) {
    LOOP(b) {
      LOOP(c) {
        ++count;
        ASSERT_NE(CgContext::GetThreadLocalContext(), nullptr);
        ASSERT_NE(CgContext::GetSharedThreadLocalContext(), nullptr);
      }
    }
  }
  ASSERT_EQ(count, 1);
  ASSERT_EQ(CgContext::GetThreadLocalContext(), nullptr);
  ASSERT_EQ(CgContext::GetSharedThreadLocalContext(), nullptr);
}
TEST(CgUtils, OptionLoopGuardContextOk) {
  Graph graph("test_graph");
  auto A = Expression::CreateExpression("A");
  auto B = Expression::CreateExpression("B");
  auto C = Expression::CreateExpression("C");
  auto a = graph.CreateAxis("a", A);
  auto b = graph.CreateAxis("b", B);
  auto c = graph.CreateAxis("c", C);

  int64_t count = 0;
  ASSERT_EQ(CgContext::GetThreadLocalContext(), nullptr);
  ASSERT_EQ(CgContext::GetSharedThreadLocalContext(), nullptr);
  LOOP(a) {
    LOOP(b) {
      OPTION_LOOP(c, LoopOption{}) {
        ++count;
        ASSERT_NE(CgContext::GetThreadLocalContext(), nullptr);
        ASSERT_EQ(CgContext::GetThreadLocalContext()->GetOption().pad_tensor_axes_to_loop, false);
        ASSERT_NE(CgContext::GetSharedThreadLocalContext(), nullptr);
      }
    }
  }
  ASSERT_EQ(count, 1);
  ASSERT_EQ(CgContext::GetThreadLocalContext(), nullptr);
  ASSERT_EQ(CgContext::GetSharedThreadLocalContext(), nullptr);

  OPTION_LOOP(a, LoopOption{.pad_tensor_axes_to_loop = true}) {
    ASSERT_NE(CgContext::GetThreadLocalContext(), nullptr);
    ASSERT_EQ(CgContext::GetThreadLocalContext()->GetOption().pad_tensor_axes_to_loop, true);
    ASSERT_NE(CgContext::GetSharedThreadLocalContext(), nullptr);
    LOOP(b) {
      LOOP(c) {
        ++count;
        ASSERT_NE(CgContext::GetThreadLocalContext(), nullptr);
        ASSERT_EQ(CgContext::GetThreadLocalContext()->GetOption().pad_tensor_axes_to_loop, false);
        ASSERT_NE(CgContext::GetSharedThreadLocalContext(), nullptr);
      }
    }
  }
  ASSERT_EQ(count, 2);
  ASSERT_EQ(CgContext::GetThreadLocalContext(), nullptr);
  ASSERT_EQ(CgContext::GetSharedThreadLocalContext(), nullptr);
}
TEST(CgUtils, NestedLoopGuardContextOk) {
  Graph graph("test_graph");
  auto A = Expression::CreateExpression("A");
  auto B = Expression::CreateExpression("B");
  auto C = Expression::CreateExpression("C");
  auto D = Expression::CreateExpression("D");
  auto a = graph.CreateAxis("a", A);
  auto b = graph.CreateAxis("b", B);
  auto c = graph.CreateAxis("c", C);
  auto d = graph.CreateAxis("d", D);

  ASSERT_EQ(CgContext::GetThreadLocalContext(), nullptr);
  ASSERT_EQ(CgContext::GetSharedThreadLocalContext(), nullptr);
  LOOP(a) {
    ASSERT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes().size(), 1UL);
    EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[0].name, a.name);
    EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[0].id, a.id);

    LOOP(b) {
      LOOP(c) {
        ASSERT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes().size(), 3UL);
        EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[0].name, a.name);
        EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[0].id, a.id);
        EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[1].name, b.name);
        EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[1].id, b.id);
        EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[2].name, c.name);
        EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[2].id, c.id);

        LOOP(d) {
          ASSERT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes().size(), 4UL);
          EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[0].name, a.name);
          EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[0].id, a.id);
          EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[1].name, b.name);
          EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[1].id, b.id);
          EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[2].name, c.name);
          EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[2].id, c.id);
          EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[3].name, d.name);
          EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[3].id, d.id);
        }

        ASSERT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes().size(), 3UL);
        EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[0].name, a.name);
        EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[0].id, a.id);
        EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[1].name, b.name);
        EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[1].id, b.id);
        EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[2].name, c.name);
        EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[2].id, c.id);
      }
    }

    ASSERT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes().size(), 1UL);
    EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[0].name, a.name);
    EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[0].id, a.id);
  }
  ASSERT_EQ(CgContext::GetThreadLocalContext(), nullptr);
  ASSERT_EQ(CgContext::GetSharedThreadLocalContext(), nullptr);
}
TEST(CgUtils, LoopGuardAxisOk) {
  Graph graph("test_graph");
  auto A = Expression::CreateExpression("A");
  auto B = Expression::CreateExpression("B");
  auto C = Expression::CreateExpression("C");
  auto a = graph.CreateAxis("a", A);
  auto b = graph.CreateAxis("b", B);
  auto c = graph.CreateAxis("c", C);

  LOOP(a) {
    LOOP(b) {
      LOOP(c) {
        ASSERT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes().size(), 3UL);
        EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[0].name, a.name);
        EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[0].id, a.id);
        EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[1].name, b.name);
        EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[1].id, b.id);
        EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[2].name, c.name);
        EXPECT_EQ(CgContext::GetThreadLocalContext()->GetLoopAxes()[2].id, c.id);
      }
    }
  }
}
TEST(CgUtils, LoopGuard_SchedAxis_Ok) {
  Graph graph("test_graph");
  auto A = Expression::CreateExpression("A");
  auto B = Expression::CreateExpression("B");
  auto C = Expression::CreateExpression("C");
  auto a = graph.CreateAxis("a", A);
  auto b = graph.CreateAxis("b", B);
  auto c = graph.CreateAxis("c", C);

  LOOP(a) {
    LOOP(b) {
      LOOP(c) {
        auto data0 = ContiguousData("data0", graph, ge::DT_FLOAT, {a, b});
        auto data1 = ContiguousData("data1", graph, ge::DT_FLOAT, {b, c});
        auto mm = MatMul("mm", data0, data1);
        (void) mm; // -Werror=unused-but-set-variable
      }
    }
  }

  auto data0 = graph.FindNode("data0");
  auto data1 = graph.FindNode("data1");
  auto mm = graph.FindNode("mm");
  ASSERT_EQ(std::vector<AxisId>(data0->attr.sched.axis), std::vector<AxisId>({a.id, b.id, c.id}));
  ASSERT_EQ(std::vector<AxisId>(data1->attr.sched.axis), std::vector<AxisId>({a.id, b.id, c.id}));
  ASSERT_EQ(std::vector<AxisId>(mm->attr.sched.axis), std::vector<AxisId>({a.id, b.id, c.id}));
}

TEST(PadTensorAxisToSched, NoContext_DoNotPad) {
  Graph graph("test_graph");
  auto A = Expression::CreateExpression("A");
  auto B = Expression::CreateExpression("B");
  auto C = Expression::CreateExpression("C");
  auto a = graph.CreateAxis("a", A);
  auto b = graph.CreateAxis("b", B);
  auto c = graph.CreateAxis("c", C);

  ge::ascir_op::Data data("data", graph);
  data.y.axis = {a.id};
  data.y.repeats = {A};
  data.y.strides = {Expression::CreateExpression(1)};

  ASSERT_TRUE(PadOutputViewToSched(data.y));
  EXPECT_EQ(data.y.axis, std::vector<AxisId>({a.id}));
  EXPECT_TRUE(*data.y.repeats[0] == *A);
  EXPECT_TRUE(*data.y.strides[0] == *Expression::CreateExpression(1));
}
TEST(PadTensorAxisToSched, NotConfigPad_DoNotPad) {
  Graph graph("test_graph");
  auto A = Expression::CreateExpression("A");
  auto B = Expression::CreateExpression("B");
  auto C = Expression::CreateExpression("C");
  auto a = graph.CreateAxis("a", A);
  auto b = graph.CreateAxis("b", B);
  auto c = graph.CreateAxis("c", C);

  ge::ascir_op::Data data("data", graph);
  data.y.axis = {a.id};
  data.y.repeats = {A};
  data.y.strides = {Expression::CreateExpression(1)};

  LOOP(a) {
    LOOP(b) {
      LOOP(c) {
        ASSERT_TRUE(PadOutputViewToSched(data.y));
      }
    }
  }
  EXPECT_EQ(data.y.axis, std::vector<AxisId>({a.id}));
  EXPECT_TRUE(*data.y.repeats[0] == *A);
  EXPECT_TRUE(*data.y.strides[0] == *Expression::CreateExpression(1));
}
TEST(PadTensorAxisToSched, NoNeedPad_Ok) {
  Graph graph("test_graph");
  auto A = Expression::CreateExpression("A");
  auto B = Expression::CreateExpression("B");
  auto C = Expression::CreateExpression("C");
  auto a = graph.CreateAxis("a", A);
  auto b = graph.CreateAxis("b", B);
  auto c = graph.CreateAxis("c", C);

  ge::ascir_op::Data data("data", graph);
  data.y.axis = {a.id, b.id, c.id};
  data.y.repeats = {A, B, C};
  data.y.strides = {C, Expression::CreateExpression(0), Expression::CreateExpression(1)};
  LOOP(a) {
    LOOP(b) {
      OPTION_LOOP(c, LoopOption{true}) {
        ASSERT_TRUE(PadOutputViewToSched(data.y));
      }
    }
  }

  EXPECT_EQ(data.y.axis, std::vector<AxisId>({a.id, b.id, c.id}));
  EXPECT_TRUE(*data.y.repeats[0] == *A);
  EXPECT_TRUE(*data.y.repeats[0] == *A);
  EXPECT_TRUE(*data.y.repeats[1] == *B);
  EXPECT_TRUE(*data.y.repeats[2] == *C);
  EXPECT_TRUE(*data.y.strides[0] == *C);
  EXPECT_TRUE(*data.y.strides[1] == *Expression::CreateExpression(0));
  EXPECT_TRUE(*data.y.strides[2] == *Expression::CreateExpression(1));
}
TEST(PadTensorAxisToSched, PadHead) {
  Graph graph("test_graph");
  auto A = Expression::CreateExpression("A");
  auto B = Expression::CreateExpression("B");
  auto C = Expression::CreateExpression("C");
  auto D = Expression::CreateExpression("d");
  auto a = graph.CreateAxis("a", A);
  auto b = graph.CreateAxis("b", B);
  auto c = graph.CreateAxis("c", C);
  auto d = graph.CreateAxis("d", D);

  ge::ascir_op::Data data("data", graph);
  data.y.SetContiguousView({c, d});
  LOOP(a) {
    LOOP(b) {
      LOOP(c) {
        OPTION_LOOP(d, LoopOption{true}) {
          ASSERT_TRUE(PadOutputViewToSched(data.y));
        }
      }
    }
  }
  EXPECT_EQ(data.y.axis, std::vector<AxisId>({a.id, b.id, c.id, d.id}));
  EXPECT_TRUE(*data.y.repeats[0] == *Expression::CreateExpression(1));
  EXPECT_TRUE(*data.y.repeats[1] == *Expression::CreateExpression(1));
  EXPECT_TRUE(*data.y.repeats[2] == *C);
  EXPECT_TRUE(*data.y.repeats[3] == *D);
  std::cout << "strides 0:" << *data.y.strides[0] << std::endl;
  std::cout << "strides 1:" << *data.y.strides[1] << std::endl;
  EXPECT_TRUE(*data.y.strides[0] == *Expression::CreateExpression(0));
  EXPECT_TRUE(*data.y.strides[1] == *Expression::CreateExpression(0));
  EXPECT_TRUE(*data.y.strides[2] == *D);
  EXPECT_TRUE(*data.y.strides[3] == *Expression::CreateExpression(1));
}
TEST(PadTensorAxisToSched, PadTail) {
  Graph graph("test_graph");
  auto A = Expression::CreateExpression("A");
  auto B = Expression::CreateExpression("B");
  auto C = Expression::CreateExpression("C");
  auto D = Expression::CreateExpression("d");
  auto a = graph.CreateAxis("a", A);
  auto b = graph.CreateAxis("b", B);
  auto c = graph.CreateAxis("c", C);
  auto d = graph.CreateAxis("d", D);

  ge::ascir_op::Data data("data", graph);
  data.y.SetContiguousView({a, b});
  LOOP(a) {
    LOOP(b) {
      LOOP(c) {
        OPTION_LOOP(d, LoopOption{true}) {
          ASSERT_TRUE(PadOutputViewToSched(data.y));
        }
      }
    }
  }

  EXPECT_EQ(data.y.axis, std::vector<AxisId>({a.id, b.id, c.id, d.id}));
  EXPECT_TRUE(*data.y.repeats[0] == *A);
  EXPECT_TRUE(*data.y.repeats[1] == *B);
  EXPECT_TRUE(*data.y.repeats[2] == *Expression::CreateExpression(1));
  EXPECT_TRUE(*data.y.repeats[3] == *Expression::CreateExpression(1));
  EXPECT_TRUE(*data.y.strides[0] == *B);
  EXPECT_TRUE(*data.y.strides[1] == *Expression::CreateExpression(1));
  EXPECT_TRUE(*data.y.strides[2] == *Expression::CreateExpression(0));
  EXPECT_TRUE(*data.y.strides[3] == *Expression::CreateExpression(0));
}
TEST(PadTensorAxisToSched, PadTail_NotContiguous) {
  Graph graph("test_graph");
  auto A = Expression::CreateExpression("A");
  auto B = Expression::CreateExpression("B");
  auto C = Expression::CreateExpression("C");
  auto D = Expression::CreateExpression("d");
  auto a = graph.CreateAxis("a", A);
  auto b = graph.CreateAxis("b", B);
  auto c = graph.CreateAxis("c", C);
  auto d = graph.CreateAxis("d", D);

  ge::ascir_op::Data data("data", graph);
  data.y.axis = {a.id, b.id, c.id};
  data.y.repeats = {A, Expression::CreateExpression(1), C};
  data.y.strides = {C, Expression::CreateExpression(0), Expression::CreateExpression(1)};

  LOOP(a) {
    LOOP(b) {
      LOOP(c) {
        OPTION_LOOP(d, LoopOption{true}) {
          ASSERT_TRUE(PadOutputViewToSched(data.y));
        }
      }
    }
  }
  EXPECT_EQ(data.y.axis, std::vector<AxisId>({a.id, b.id, c.id, d.id}));
  EXPECT_TRUE(*data.y.repeats[0] == *A);
  EXPECT_TRUE(*data.y.repeats[1] == *Expression::CreateExpression(1));
  EXPECT_TRUE(*data.y.repeats[2] == *C);
  EXPECT_TRUE(*data.y.repeats[3] == *Expression::CreateExpression(1));
  EXPECT_TRUE(*data.y.strides[0] == *C);
  EXPECT_TRUE(*data.y.strides[1] == *Expression::CreateExpression(0));
  EXPECT_TRUE(*data.y.strides[2] == *Expression::CreateExpression(1));
  EXPECT_TRUE(*data.y.strides[3] == *Expression::CreateExpression(0));
}
TEST(PadTensorAxisToSched, PadMiddle) {
  Graph graph("test_graph");
  auto A = Expression::CreateExpression("A");
  auto B = Expression::CreateExpression("B");
  auto C = Expression::CreateExpression("C");
  auto D = Expression::CreateExpression("d");
  auto a = graph.CreateAxis("a", A);
  auto b = graph.CreateAxis("b", B);
  auto c = graph.CreateAxis("c", C);
  auto d = graph.CreateAxis("d", D);

  ge::ascir_op::Data data("data", graph);
  data.y.SetContiguousView({a, d});

  LOOP(a) {
    LOOP(b) {
      LOOP(c) {
        OPTION_LOOP(d, LoopOption{true}) {
          ASSERT_TRUE(PadOutputViewToSched(data.y));
        }
      }
    }
  }
  EXPECT_EQ(data.y.axis, std::vector<AxisId>({a.id, b.id, c.id, d.id}));
  EXPECT_TRUE(*data.y.repeats[0] == *A);
  EXPECT_TRUE(*data.y.repeats[1] == *Expression::CreateExpression(1));
  EXPECT_TRUE(*data.y.repeats[2] == *Expression::CreateExpression(1));
  EXPECT_TRUE(*data.y.repeats[3] == *D);
  EXPECT_TRUE(*data.y.strides[0] == *D);
  EXPECT_TRUE(*data.y.strides[1] == *Expression::CreateExpression(0));
  EXPECT_TRUE(*data.y.strides[2] == *Expression::CreateExpression(0));
  EXPECT_TRUE(*data.y.strides[3] == *Expression::CreateExpression(1));
}
TEST(PadTensorAxisToSched, PadMultiple) {
  Graph graph("test_graph");
  auto A = Expression::CreateExpression("A");
  auto B = Expression::CreateExpression("B");
  auto C = Expression::CreateExpression("C");
  auto D = Expression::CreateExpression("d");
  auto a = graph.CreateAxis("a", A);
  auto b = graph.CreateAxis("b", B);
  auto c = graph.CreateAxis("c", C);
  auto d = graph.CreateAxis("d", D);

  ge::ascir_op::Data data("data", graph);
  data.y.SetContiguousView({b, d});

  LOOP(a) {
    LOOP(b) {
      LOOP(c) {
        OPTION_LOOP(d, LoopOption{true}) {
          ASSERT_TRUE(PadOutputViewToSched(data.y));
        }
      }
    }
  }
  EXPECT_EQ(data.y.axis, std::vector<AxisId>({a.id, b.id, c.id, d.id}));
  EXPECT_EQ(data.y.axis, std::vector<AxisId>({a.id, b.id, c.id, d.id}));
  EXPECT_TRUE(*data.y.repeats[0] == *Expression::CreateExpression(1));
  EXPECT_TRUE(*data.y.repeats[1] == *B);
  EXPECT_TRUE(*data.y.repeats[2] == *Expression::CreateExpression(1));
  EXPECT_TRUE(*data.y.repeats[3] == *D);
  EXPECT_TRUE(*data.y.strides[0] == *Expression::CreateExpression(0));
  EXPECT_TRUE(*data.y.strides[1] == *D);
  EXPECT_TRUE(*data.y.strides[2] == *Expression::CreateExpression(0));
  EXPECT_TRUE(*data.y.strides[3] == *Expression::CreateExpression(1));
}
TEST(PadTensorAxisToSched, SameAxisNumButNotMatch_Failed) {
  Graph graph("test_graph");
  auto A = Expression::CreateExpression("A");
  auto B = Expression::CreateExpression("B");
  auto C = Expression::CreateExpression("C");
  auto D = Expression::CreateExpression("d");
  auto a = graph.CreateAxis("a", A);
  auto b = graph.CreateAxis("b", B);
  auto c = graph.CreateAxis("c", C);
  auto d = graph.CreateAxis("d", D);

  ge::ascir_op::Data data("data", graph);
  data.y.SetContiguousView({b, a, c, d});

  LOOP(a) {
    LOOP(b) {
      LOOP(c) {
        OPTION_LOOP(d, LoopOption{true}) {
          ASSERT_FALSE(PadOutputViewToSched(data.y));
        }
      }
    }
  }
}
TEST(PadTensorAxisToSched, DiffAxisNumAndNotMatch1_Failed) {
  Graph graph("test_graph");
  auto A = Expression::CreateExpression("A");
  auto B = Expression::CreateExpression("B");
  auto C = Expression::CreateExpression("C");
  auto D = Expression::CreateExpression("d");
  auto a = graph.CreateAxis("a", A);
  auto b = graph.CreateAxis("b", B);
  auto c = graph.CreateAxis("c", C);
  auto d = graph.CreateAxis("d", D);

  ge::ascir_op::Data data("data", graph);
  data.y.SetContiguousView({a, b, c, d});

  LOOP(a) {
    LOOP(b) {
      OPTION_LOOP(c, LoopOption{true}) {
        ASSERT_FALSE(PadOutputViewToSched(data.y));
      }
    }
  }
}
TEST(PadTensorAxisToSched, DiffAxisNumAndNotMatch2_Failed) {
  Graph graph("test_graph");
  auto A = Expression::CreateExpression("A");
  auto B = Expression::CreateExpression("B");
  auto C = Expression::CreateExpression("C");
  auto D = Expression::CreateExpression("d");
  auto a = graph.CreateAxis("a", A);
  auto b = graph.CreateAxis("b", B);
  auto c = graph.CreateAxis("c", C);
  auto d = graph.CreateAxis("d", D);

  ge::ascir_op::Data data("data", graph);
  data.y.SetContiguousView({a, c, b});

  LOOP(a) {
    LOOP(b) {
      LOOP(c) {
        OPTION_LOOP(d, LoopOption{true}) {
          ASSERT_FALSE(PadOutputViewToSched(data.y));
        }
      }
    }
  }
}
TEST(AutoPadAxis, Ok) {
  Graph graph("test_graph");
  auto A = Expression::CreateExpression("A");
  auto B = Expression::CreateExpression("B");
  auto C = Expression::CreateExpression("C");
  auto D = Expression::CreateExpression("d");
  auto a = graph.CreateAxis("a", A);
  auto b = graph.CreateAxis("b", B);
  auto c = graph.CreateAxis("c", C);
  auto d = graph.CreateAxis("d", D);

  LOOP(a) {
    LOOP(b) {
      LOOP(c) {
        OPTION_LOOP(d, LoopOption{true}) {
          auto data0 = ContiguousData("data0", graph, ge::DT_FLOAT16, {a, b, d});
          auto data1 = ContiguousData("data1", graph, ge::DT_FLOAT16, {a, c, d});
          AscendString name;
          data1.GetOwnerOp().GetName(name);
          EXPECT_EQ("data1", std::string(name.GetString()));
          auto load0 = Load("load0", data0);
          auto load1 = Load("load1", data0);
          auto mm = MatMul("mm", load0, load1);
          mm.SetContiguousView({a, b, c});
          PadOutputViewToSched(mm);
          auto data2 = ContiguousData("data1", graph, ge::DT_FLOAT, {a, c, d});
          auto y = CalcY("y", data0, data2, data1, mm);
          EXPECT_EQ(y.dtype, ge::DT_FLOAT);
        }
      }
    }
  }

  auto d0 = graph.FindNode("data0");
  EXPECT_EQ(d0->outputs[0].axis, std::vector<AxisId>({a.id, b.id, c.id, d.id}));
  EXPECT_TRUE(*d0->outputs[0].repeats[0] == *A);
  EXPECT_TRUE(*d0->outputs[0].repeats[1] == *B);
  EXPECT_TRUE(*d0->outputs[0].repeats[2] == *Expression::CreateExpression(1));
  EXPECT_TRUE(*d0->outputs[0].repeats[3] == *D);
  EXPECT_TRUE(*d0->outputs[0].strides[0] == *(B*D));
  EXPECT_TRUE(*d0->outputs[0].strides[1] == *D);
  EXPECT_TRUE(*d0->outputs[0].strides[2] == *Expression::CreateExpression(0));
  EXPECT_TRUE(*d0->outputs[0].strides[3] == *Expression::CreateExpression(1));

  auto d1 = graph.FindNode("data1");

  EXPECT_EQ(d0->outputs[0].axis, std::vector<AxisId>({a.id, b.id, c.id, d.id}));
  EXPECT_TRUE(*d1->outputs[0].repeats[0] == *A);
  EXPECT_TRUE(*d1->outputs[0].repeats[1] == *Expression::CreateExpression(1));
  EXPECT_TRUE(*d1->outputs[0].repeats[2] == *C);
  EXPECT_TRUE(*d1->outputs[0].repeats[3] == *D);
  EXPECT_TRUE(*d1->outputs[0].strides[0] == *(C*D));
  EXPECT_TRUE(*d1->outputs[0].strides[1] == *Expression::CreateExpression(0));
  EXPECT_TRUE(*d1->outputs[0].strides[2] == *D);
  EXPECT_TRUE(*d1->outputs[0].strides[3] == *Expression::CreateExpression(1));

  auto mm = graph.FindNode("mm");
  EXPECT_EQ(mm->outputs[0].axis, std::vector<AxisId>({a.id, b.id, c.id, d.id}));
  EXPECT_TRUE(*mm->outputs[0].repeats[0] == *A);
  EXPECT_TRUE(*mm->outputs[0].repeats[1] == *B);
  EXPECT_TRUE(*mm->outputs[0].repeats[2] == *C);
  EXPECT_TRUE(*mm->outputs[0].repeats[3] == *Expression::CreateExpression(1));
  EXPECT_TRUE(*mm->outputs[0].strides[0] == *(B*C));
  EXPECT_TRUE(*mm->outputs[0].strides[1] == *C);
  EXPECT_TRUE(*mm->outputs[0].strides[2] == *Expression::CreateExpression(1));
  EXPECT_TRUE(*mm->outputs[0].strides[3] == *Expression::CreateExpression(0));
}

TEST(SetDataNodeAttr, Ok) {
  Graph graph("test_graph");
  auto A = Expression::CreateExpression("A");
  auto B = Expression::CreateExpression("B");
  auto C = Expression::CreateExpression("C");
  auto D = Expression::CreateExpression("d");
  auto a = graph.CreateAxis("a", A);
  auto b = graph.CreateAxis("b", B);
  auto c = graph.CreateAxis("c", C);
  auto d = graph.CreateAxis("d", D);

  LOOP(a) {
    LOOP(b) {
      LOOP(c) {
        OPTION_LOOP(d, LoopOption{true}) {
          auto data0 = ContiguousData("data0", graph, ge::DT_FLOAT16, {a, b, d}, 0, {{1}, {2}}, ge::FORMAT_CN);
          auto data1 = ContiguousData("data1", graph, ge::DT_FLOAT16, {a, c, d}, 1, {{1}, {2}, {3}}, ge::FORMAT_CN);
          AscendString name;
          data0.GetOwnerOp().GetName(name);
          EXPECT_EQ("data0", std::string(name.GetString()));
          data1.GetOwnerOp().GetName(name);
          EXPECT_EQ("data1", std::string(name.GetString()));
        }
      }
    }
  }

  auto d0 = graph.FindNode("data0");
  ge::GeAttrValue attr_value;
  int64_t index_value = -1;
  (void)d0->GetOpDesc()->GetAttr("index", attr_value);
  attr_value.GetValue<int64_t>(index_value);
  EXPECT_EQ(index_value, 0);

  std::vector<std::vector<int64_t>> axis_continuous_map;
  (void)d0->GetOpDesc()->GetAttr("axis_continuous_map", attr_value);
  attr_value.GetValue<std::vector<std::vector<int64_t>>>(axis_continuous_map);
  EXPECT_EQ(axis_continuous_map.size(), 2);

  int64_t format_value = -1;
  (void)d0->GetOpDesc()->GetAttr("format", attr_value);
  attr_value.GetValue<int64_t>(format_value);
  EXPECT_EQ(format_value, static_cast<int64_t>(ge::FORMAT_CN));

  auto d1 = graph.FindNode("data1");
  index_value = -1;
  (void)d1->GetOpDesc()->GetAttr("index", attr_value);
  attr_value.GetValue<int64_t>(index_value);
  EXPECT_EQ(index_value, 1);

  axis_continuous_map.clear();
  (void)d1->GetOpDesc()->GetAttr("axis_continuous_map", attr_value);
  attr_value.GetValue<std::vector<std::vector<int64_t>>>(axis_continuous_map);
  EXPECT_EQ(axis_continuous_map.size(), 3);

  format_value = -1;
  (void)d1->GetOpDesc()->GetAttr("format", attr_value);
  attr_value.GetValue<int64_t>(format_value);
  EXPECT_EQ(format_value, static_cast<int64_t>(ge::FORMAT_CN));
}
}  // namespace cg
}  // namespace ascir
}
